/*	Interactive Gravitational Simulator
	Project: IGS Visualization
	File: Visualizer.cpp
	Author: Mike Bantegui <mbante2@gmail.com>, Hofstra University 
	Copyright (C) 2012 - 2013 Mike Bantegui

	This file is part of the IGS software package for simulating
	the N-Body problem in real-time. This was developed as part
	of a departmental honors thesis project at Hofstra University
	in the Spring 2012 semester.

	IGS is free software; you can redistribute it and/or modify
	it under the terms of the GNU General Public License as published by
	the Free Software Foundation; either version 2 of the License, or
	(at your option) any later version.

	IGS is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU General Public License for more details.

	You should have received a copy of the GNU General Public License
	along with this program; if not, write to the Free Software
	Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.
*/

#include "Body.h"
#include "Integrator.h"
#include "Model.h"
#include "NBodySystem.h"
#include "Tree.h"
#include "Simulator.h"
#include "Snapshot.h"
#include "Visualizer.h"
#include "ViewNode.h"
#include <iostream>
#include <SDL_opengl.h>

namespace IGS
{

void DrawBox();
void DrawTree(std::shared_ptr<ViewNode> root, unsigned depth, unsigned &drawn, bool iterative);
void DrawAxis(bool withLabel);
Vector CameraVector(double theta, double phi);
Vector UpVector(double theta, double phi);
Vector SideVector(double theta, double phi);

Visualizer::Visualizer() : UseFreeCamera(false), Position(0, 0, -15),
	Distance(15), Theta(0), Phi(0), MouseX(0), MouseY(0), Fonts(),
	IsPaused(false), IsUpdating(true), TreeVisible(false), IsIterativeTree(false),
	MaxNodes(0), _Simulator(NULL), Dt(0.001), Generation(0), Snapshots(),
	MaxSnapshots(32), SnapshotIndex(0)
{
}

bool Visualizer::Initialize(int argc, char **argv)
{
	glClearColor(0, 0, 0, 0);
	glDisable(GL_DEPTH_TEST);
	glEnable(GL_POINT_SMOOTH);
	glEnable(GL_LINE_SMOOTH);
	glEnable(GL_TEXTURE_2D);
	glPointSize(2.5f);

	glViewport(0, 0, Width(), Height());

	if (!Fonts.LoadFont("C:\\Windows\\Fonts\\tahoma.ttf", 12))
		return false;

	int N = 512;
	unsigned long long MaxMemory = (3ULL * 1024) * (1024 * 512); // 1.5 GB
	double theta = 2;
	bool useOctree = false;

	if (argc > 1)
		N = atoi(argv[1]);
	if (argc > 2)
		useOctree = !!atoi(argv[2]);
	if (argc > 3 && useOctree)
		theta = atof(argv[3]);
	if (argc > 4)
		MaxMemory = atol(argv[3]);

	double softening = sqrt(1.0 / N);
	std::cerr << "Using softening length e = " << softening << std::endl;

	MaxSnapshots = MaxMemory / (N * sizeof(PhasePoint));
	std::cerr << "Maxmimum of " << MaxSnapshots;
	std::cerr << " snapshots will be taken" << std::endl;

	std::shared_ptr<Model> model(new PlummerModel());
/*	std::shared_ptr<CompositeModel> model(new CompositeModel());
	model->AddComponent(new DiskModel(), Vector(-8, 0, 0), Vector(+0.1, +0.3, 0), 512);
	model->AddComponent(new DiskModel(), Vector(+8, 0, 0), Vector(-0.1, -0.3, 0), 512);*/

	NBodySystem *system = model->Create(N, true);
	system->SetSoftening(softening);
	
	std::string integratorName = "hermite-6";
	std::string treeName = (useOctree ? "Octree" : "None");

	Integrator *integrator = CreateIntegrator(integratorName);
	Tree *tree = CreateTree(treeName, theta);
#ifdef PoolAllocator
	if (useOctree)
		tree->SetPoolSize(N);
#endif

	Simulator *sim = new Simulator(system, integrator, StepShared, tree);
	_Simulator = std::shared_ptr<Simulator>(sim);

	return true;
}

void Visualizer::Cleanup()
{
	
}

void Visualizer::Draw()
{
	glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT);

	Set3DContext();
	SetCamera();

	DrawParticles();
	DrawAxis(true);
	if (TreeVisible)
		DrawTree();

	Set2DContext();
	DrawOSD();

	SDL_GL_SwapBuffers();

	if (!IsUpdating && !IsPaused)
	{
		if (TreeVisible && IsIterativeTree)
		{
			unsigned maxCount = Snapshots[SnapshotIndex]->Tree()->NodeCount;
			MaxNodes = std::min(MaxNodes + 1, maxCount);
		}

		else
			SnapshotIndex = (SnapshotIndex + 1) % Snapshots.size();
	}
}

void Visualizer::Keys()
{
	for (unsigned i = 0; i < KeyEvents.size(); i++)
	{
		const SDL_Event &event = KeyEvents[i];
		SDLKey key = event.key.keysym.sym;

		if (event.type == SDL_KEYDOWN)
		{
			switch (key)
			{
			case SDLK_f:
				UseFreeCamera = !UseFreeCamera;
				break;
			case SDLK_r:
				if (Snapshots.size() > 0)
					IsUpdating = !IsUpdating;
				IsIterativeTree = false;
				break;
			case SDLK_p:
				IsPaused = !IsPaused;
				break;
			case SDLK_t:
				TreeVisible = !TreeVisible;
				IsIterativeTree = false;
				break;
			case SDLK_i:
				if (!IsUpdating && TreeVisible)
				{
					IsIterativeTree = !IsIterativeTree;
					MaxNodes = 0;
				}
				break;
			default:
				KeyState[key] = true;
			}
		}

		else
			KeyState[key] = false;
	}

	if (KeyState[SDLK_ESCAPE])
	{
		Quit();
		return;
	}

	if (UseFreeCamera)
	{
		if (KeyState[SDLK_w])
			Position += 0.1 * CameraVector(Theta, Phi);
		else if (KeyState[SDLK_s])
			Position -= 0.1 * CameraVector(Theta, Phi);
		if (KeyState[SDLK_a])
			Position += 0.1 * SideVector(Theta, Phi);
		else if (KeyState[SDLK_d])
			Position -= 0.1 * SideVector(Theta, Phi);
	}

	else
	{
		if (KeyState[SDLK_w])
			Distance = std::max(0.05, Distance - 0.05);
		else if (KeyState[SDLK_s])
			Distance += 0.05;

	}
}

void Visualizer::Mouse()
{
	int oldMouseX = MouseX;
	int oldMouseY = MouseY;

	int button = SDL_GetMouseState(&MouseX, &MouseY);

	if (button != 0)
	{
		Theta += 0.01 * (MouseY - oldMouseY);
		Phi += 0.01 * (MouseX - oldMouseX);
	}
}

void Visualizer::Loop()
{
	if (!IsUpdating || IsPaused)
	{
		SDL_Delay(25);
		return;
	}

	_Simulator->Evolve(Dt);

	// Periodically take a snapshot
	if (Generation % 5 == 0 && Snapshots.size() < MaxSnapshots)
	{
		const NBodySystem &system = *_Simulator->GetSystem();
		Snapshot *snapshot = system.ToSnapshot();
		Snapshots.push_back(std::shared_ptr<Snapshot>(snapshot));
	}

	Generation++;
}

void Visualizer::SetCamera()
{
	static const double RadToDeg = 90 / asin(1.0);

	if (UseFreeCamera)
	{
		glRotated(Theta * RadToDeg, 1, 0, 0);
		glRotated(Phi * RadToDeg, 0, 1, 0);
		glTranslated(Position.x, Position.y, Position.z);
	}

	else
	{
		const NBodySystem &system = *_Simulator->GetSystem();

		Vector center;
		if (IsUpdating)
			center = system.CenterOfMass();
		else
			center = Snapshots[SnapshotIndex]->CenterOfMass();

		Vector r = CameraVector(Theta, Phi);
		Vector up = UpVector(Theta, Phi);
		Vector eye = center + Distance * r;

		gluLookAt(eye.x, eye.y, eye.z,
			center.x, center.y, center.z,
			up.x, up.y, up.z);
	}
}

template <class TSystem, class TBody>
void RenderParticles(const TSystem &system)
{
	const double &maxV = system.MaxVelocity();

	glBegin(GL_POINTS);
	for (int i = 0; i < system.Size(); i++)
	{
		const TBody &body = system[i];

		const double &v = Length(body.Velocity);
		const double &color = 1 - v / maxV;
		glColor4d(1, color, 0, 0.7);

		const Vector &r = body.Position;
		glVertex3d(r.x, r.y, r.z);
	}
	glEnd();
}

void Visualizer::DrawParticles()
{
	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

	if (IsUpdating)
	{
		const NBodySystem &system = *_Simulator->GetSystem();
		RenderParticles<NBodySystem, Body>(system);
	}

	else
	{
		const Snapshot &snapshot = *Snapshots[SnapshotIndex];
		RenderParticles<Snapshot, PhasePoint>(snapshot);

		int prev = std::max(((int)SnapshotIndex) - 5, 0);

		// Draw tail to depict the trajectory
		if (prev != SnapshotIndex)
		{
			const Snapshot &prevSnapshot = *Snapshots[SnapshotIndex];
			double maxV = prevSnapshot.MaxVelocity();

			glBegin(GL_LINES);
			for (int i = 0; i < snapshot.Size(); i++)
			{
				const PhasePoint &b1 = snapshot[i];
				const PhasePoint &b2 = prevSnapshot[i];

				const Vector &r1 = b1.Position;
				const Vector &r2 = b2.Position;

				const double v = Length(b1.Velocity);
				double color = 1 - v / maxV;

				glColor4d(1, color, 0, 0.7);
				glVertex3d(r2.x, r2.y, r2.z);
				glVertex3d(r1.x, r1.y, r1.z);
			}
			glEnd();
		}
	}

	glDisable(GL_BLEND);
}

double KE, PE, TE;
Vector MC, PC, L, w, F, minB, maxB;

void Visualizer::DrawOSD()
{
	double averageFPS = 1000 * DrawTimer.AverageFPS(); // Rescale to frames per second
	double averageFrameTime = DrawTimer.AverageFrameTime();
	int Step = _Simulator->TotalSteps();

	if (IsUpdating && DrawTimer.TotalSamples() % 5 == 0)
	{
		const NBodySystem &system = *_Simulator->GetSystem();
		KE = system.KineticEnergy();
		PE = system.PotentialEnergy();
		TE = KE + PE;
		MC = system.CenterOfMass();
		PC = system.CenterOfMomentum();
		L = system.AngularMomentum();
		w = system.AngularVelocity();
		F = system.TotalForce();
		minB = system.MinBound();
		maxB = system.MaxBound();
	}

	else if (!IsUpdating)
	{
		const Snapshot &snapshot = *Snapshots[SnapshotIndex];

		KE = snapshot.KineticEnergy();
		PE = snapshot.PotentialEnergy();
		TE = KE + PE;
		MC = snapshot.CenterOfMass();
		PC = snapshot.CenterOfMomentum();
		L = snapshot.AngularMomentum();
		w = snapshot.AngularVelocity();
		F = snapshot.TotalForce();
		minB = snapshot.MinBound();
		maxB = snapshot.MaxBound();
	}

	const std::string labels[] = { "Average FPS", "Average Frame Time",
		"Kinetic Energy", "Potential Energy", "Total Energy", "Center of Mass",
		"Center of Momentum", "Angular Momentum", "Angular Velocity", "Force", "Min Boundary",
		"Max Boundary", "Bodies", "Step #", "Snapshot #", "Integration Steps", "t" };
	const std::string values[] = { ToString(averageFPS), ToString(averageFrameTime),
		ToString(KE), ToString(PE), ToString(TE), ToString(MC), ToString(PC),
		ToString(L), ToString(w), ToString(F), ToString(minB), ToString(maxB), 
		ToString(_Simulator->GetSystem()->Size()), ToString(Step), ToString(SnapshotIndex),
		ToString(_Simulator->Steps()), ToString(_Simulator->Time()) };
	int ns = 17;

	int lwidth = 0;
	for (int i = 0; i < ns; i++)
		lwidth = std::max(lwidth, Fonts.WordWidth(labels[i]));

	int y = 10, dy = Fonts.FontHeight();
	int x = 10, dx = lwidth + Fonts.WordWidth("M");

	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);

	glColor4d(1, 1, 1, 0.8);
	for (int i = 0; i < ns; i++)
		Fonts.Render(x, y + i * dy, labels[i]);

	// Linearly interpolate between green, yellow and red
	// colors for the frame rate
	if (averageFPS > 30)
		glColor4d((60 - averageFPS) / 30, 1, 0, 0.9);
	else if (averageFPS > 20)
		glColor4d(1, (averageFPS - 20) / 10, 0, 0.9);
	else
		glColor4d(1, 0, 0, 0.9);
	for (int i = 0; i < 2; i++)
		Fonts.Render(x + dx, y + i * dy, values[i]);

	glColor4d(1, 1, 1, 1);
	for (int i = 2; i < ns; i++)
		Fonts.Render(x + dx, y + i * dy, values[i]);

	glDisable(GL_BLEND);
}

void Visualizer::DrawTree()
{
	glEnable(GL_BLEND);
	glPolygonMode(GL_FRONT_AND_BACK, GL_LINE);

	unsigned drawn = 0;

	std::shared_ptr<ViewNode> root;
	
	if (IsUpdating)
		root = _Simulator->GetSystem()->Tree();
	else
		root = Snapshots[SnapshotIndex]->Tree();
	DrawTree(root, drawn, 1);

	glPolygonMode(GL_FRONT_AND_BACK, GL_FILL);
	glDisable(GL_BLEND);
}

void Visualizer::DrawTree(std::shared_ptr<ViewNode> root, unsigned &nodeCount, unsigned depth)
{
	if (root == NULL)
		return;

	if (IsIterativeTree && nodeCount == MaxNodes)
		return;

	nodeCount++;

	const Vector &r = root->Position;
	const double size = 0.805996 * root->Radius;
	std::vector<std::shared_ptr<ViewNode>> nodes = root->Branches;
	
	for (unsigned i = 0; i < nodes.size(); i++)
		DrawTree(nodes[i], nodeCount, depth + 1);

	// Color based on depth
	// 0.0484 is an arbitrary factor which happens to give good results
	double v = 1 - depth / (double) root->Depth;
	double visibility = 0.0484 * v * v;

	glPushMatrix();
	glTranslated(r.x, r.y, r.z);
	glScaled(size, size, size);
	glColor4d(0, 0, 1, visibility);

	DrawBox();
	glPopMatrix();
}

void DrawBox()
{
	glBegin(GL_QUADS);

	// +Z
	glVertex3i(+1, +1, +1); glVertex3i(+1, -1, +1);
	glVertex3i(-1, -1, +1); glVertex3i(-1, +1, +1);
	// -Z
	glVertex3i(+1, +1, -1); glVertex3i(+1, -1, -1);
	glVertex3i(-1, -1, -1); glVertex3i(-1, +1, -1);
	// +X
	glVertex3i(+1, +1, +1); glVertex3i(+1, -1, +1);
	glVertex3i(+1, -1, -1); glVertex3i(+1, +1, -1);
	// -X
	glVertex3i(-1, +1, +1); glVertex3i(-1, -1, +1);
	glVertex3i(-1, -1, -1); glVertex3i(-1, +1, -1);
	// +Y
	glVertex3i(-1, +1, -1); glVertex3i(-1, +1, +1);
	glVertex3i(+1, +1, +1); glVertex3i(+1, +1, -1);
	// -Y
	glVertex3i(-1, -1, -1); glVertex3i(-1, -1, +1);
	glVertex3i(+1, -1, +1); glVertex3i(+1, -1, -1);

	glEnd();
}

void DrawAxis(bool withLabel)
{
	glEnable(GL_BLEND);
	glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA);
	glBegin(GL_LINES);

		// X-axis line
		glColor4d(1, 0, 0, 0.7);
		glVertex3d(-10, 0, 0); glVertex3d(+10, 0, 0);
		// Y-axis line
		glColor4d(0, 1, 0, 0.7);
		glVertex3d(0, -10, 0); glVertex3d(0, +10, 0);
		// Z-axis line
		glColor4d(0, 0, 1, 0.7);
		glVertex3d(0, 0, -10); glVertex3d(0, 0, +10);

	if (withLabel)
	{
		// "X" polygon
		glVertex3d(+11, -0.5, -0.5); glVertex3d(+11, +0.5, +0.5);
		glVertex3d(+11, -0.5, +0.5); glVertex3d(+11, +0.5, -0.5);
		// "Y" polygon
		glVertex3d(0, +11, -0.25); glVertex3d(0, +11, +0.25);
		glVertex3d(0, +11, +0.25); glVertex3d(-0.5, +11, +0.75);
		glVertex3d(0, +11, +0.25); glVertex3d(+0.5, +11, +0.75);
		// "Z" polygon
		glVertex3d(-0.5, +0.5, +11); glVertex3d(+0.5, +0.5, +11);
		glVertex3d(+0.5, +0.5, +11); glVertex3d(-0.5, -0.5, +11);
		glVertex3d(-0.5, -0.5, +11); glVertex3d(+0.5, -0.5, +11);
	}
	glEnd();
	glDisable(GL_BLEND);
}

Vector CameraVector(double theta, double phi)
{
	return Vector(-cos(theta) * sin(phi),
		sin(theta), cos(theta) * cos(phi));
}

Vector UpVector(double theta, double phi)
{
	return Vector(sin(theta) * sin(phi),
		cos(theta), -sin(theta) * cos(phi));
}

Vector SideVector(double theta, double phi)
{
	return Vector(cos(theta) * cos(phi),
		0, cos(theta) * sin(phi));
}

}